{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Number of Good Leaf Nodes Pairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Medium"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #tree #depth-first-search #binary-tree"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #树 #深度优先搜索 #二叉树"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: countPairs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #好叶子节点对的数量"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你二叉树的根节点 <code>root</code> 和一个整数 <code>distance</code> 。</p>\n",
    "\n",
    "<p>如果二叉树中两个 <strong>叶</strong> 节点之间的 <strong>最短路径长度</strong> 小于或者等于 <code>distance</code> ，那它们就可以构成一组 <strong>好叶子节点对</strong> 。</p>\n",
    "\n",
    "<p>返回树中 <strong>好叶子节点对的数量</strong> 。</p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode-cn.com/aliyun-lc-upload/uploads/2020/07/26/e1.jpg\" style=\"height: 321px; width: 321px;\"></p>\n",
    "\n",
    "<pre><strong>输入：</strong>root = [1,2,3,null,4], distance = 3\n",
    "<strong>输出：</strong>1\n",
    "<strong>解释：</strong>树的叶节点是 3 和 4 ，它们之间的最短路径的长度是 3 。这是唯一的好叶子节点对。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "\n",
    "<p><img alt=\"\" src=\"https://assets.leetcode-cn.com/aliyun-lc-upload/uploads/2020/07/26/e2.jpg\" style=\"height: 321px; width: 441px;\"></p>\n",
    "\n",
    "<pre><strong>输入：</strong>root = [1,2,3,4,5,6,7], distance = 3\n",
    "<strong>输出：</strong>2\n",
    "<strong>解释：</strong>好叶子节点对为 [4,5] 和 [6,7] ，最短路径长度都是 2 。但是叶子节点对 [4,6] 不满足要求，因为它们之间的最短路径长度为 4 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 3：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>root = [7,1,4,6,null,5,3,null,null,null,null,null,2], distance = 3\n",
    "<strong>输出：</strong>1\n",
    "<strong>解释：</strong>唯一的好叶子节点对是 [2,5] 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 4：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>root = [100], distance = 1\n",
    "<strong>输出：</strong>0\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 5：</strong></p>\n",
    "\n",
    "<pre><strong>输入：</strong>root = [1,1,1], distance = 2\n",
    "<strong>输出：</strong>1\n",
    "</pre>\n",
    "\n",
    "<p>&nbsp;</p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>tree</code> 的节点数在 <code>[1, 2^10]</code> 范围内。</li>\n",
    "\t<li>每个节点的值都在 <code>[1, 100]</code> 之间。</li>\n",
    "\t<li><code>1 &lt;= distance &lt;= 10</code></li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [number-of-good-leaf-nodes-pairs](https://leetcode.cn/problems/number-of-good-leaf-nodes-pairs/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [number-of-good-leaf-nodes-pairs](https://leetcode.cn/problems/number-of-good-leaf-nodes-pairs/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[1,2,3,null,4]\\n3', '[1,2,3,4,5,6,7]\\n3', '[7,1,4,6,null,5,3,null,null,null,null,null,2]\\n3']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        res=[]\n",
    "        temp=[root]\n",
    "        ss=[[0]]\n",
    "        while temp and ss:\n",
    "            s=[]\n",
    "            new=[]\n",
    "            for i in range(len(temp)):\n",
    "                if temp[i].left:\n",
    "                    new.append(temp[i].left)\n",
    "                    s.append(ss[i]+[1])\n",
    "                if temp[i].right:\n",
    "                    new.append(temp[i].right)\n",
    "                    s.append(ss[i]+[-1])\n",
    "                if not temp[i].left and not temp[i].right:\n",
    "                    res.append(ss[i])\n",
    "            temp=new \n",
    "            ss=s\n",
    "        def finddis(l1,l2):\n",
    "            ll1=len(l1)\n",
    "            ll2=len(l2)\n",
    "            i=0\n",
    "            j=0\n",
    "            while i<ll1 and j<ll2:\n",
    "                if l1[i]==l2[j]:\n",
    "                    i+=1\n",
    "                    j+=1\n",
    "                else:\n",
    "                    return ll1-i+ll2-j\n",
    "        num=0\n",
    "        for i in range(len(res)):\n",
    "            for j in range(i+1,len(res)):\n",
    "                if finddis(res[i], res[j])<=distance:\n",
    "                    num+=1\n",
    "        return num\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        graph = collections.defaultdict(list)\n",
    "        leaves = set()\n",
    "        stk = [root]\n",
    "        while stk:  # 先遍历建图，并找到所有叶子节点\n",
    "            node = stk.pop()\n",
    "            if not node.left and not node.right:\n",
    "                leaves.add(node)\n",
    "            if node.left:\n",
    "                graph[node.left].append(node)\n",
    "                graph[node].append(node.left)\n",
    "                stk.append(node.left)\n",
    "            if node.right:\n",
    "                graph[node.right].append(node)\n",
    "                graph[node].append(node.right)\n",
    "                stk.append(node.right)\n",
    "        count = 0 # valid leaves count\n",
    "        for leaf in leaves:\n",
    "            dq = collections.deque([(leaf, 0)]) # 从每一个叶子出发，去找<=distance的另一个叶子\n",
    "            seen = set([leaf])\n",
    "            while dq:\n",
    "                node, dist = dq.popleft()\n",
    "                for nei in graph[node]:\n",
    "                    if nei not in seen and dist + 1 <= distance:\n",
    "                        seen.add(nei)\n",
    "                        dq.append((nei, dist + 1))\n",
    "                    if node != leaf and node in leaves and dist <= distance:\n",
    "                        count += 1 # find another leaf within the distance\n",
    "        return count // 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def dfs(self, node: TreeNode):\n",
    "        if not node.left and not node.right:\n",
    "            return [0]\n",
    "        if node.left:\n",
    "            left_leaf = self.dfs(node.left)\n",
    "            left_leaf = [ll+1 for ll in left_leaf if ll<self.distance]\n",
    "        else:\n",
    "            left_leaf = []\n",
    "        if node.right:\n",
    "            right_leaf = self.dfs(node.right)\n",
    "            right_leaf = [rl+1 for rl in right_leaf if rl<self.distance]\n",
    "        else:\n",
    "            right_leaf = []\n",
    "        if node.left and node.right:\n",
    "            print(node.val,left_leaf,right_leaf)\n",
    "            for ll in left_leaf:\n",
    "                for rl in right_leaf:\n",
    "                    if ll + rl <= self.distance:\n",
    "                        self.rs += 1\n",
    "\n",
    "        return left_leaf + right_leaf\n",
    "\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.rs = 0\n",
    "        self.distance = distance\n",
    "        self.dfs(root)\n",
    "        return self.rs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.count = 0\n",
    "        \n",
    "        # DFS 返回每个节点的所有叶子节点的深度\n",
    "        def dfs(node):\n",
    "            if not node:\n",
    "                return []\n",
    "            \n",
    "            # 如果是叶子节点，返回深度为1\n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            \n",
    "            left_depths = dfs(node.left)\n",
    "            right_depths = dfs(node.right)\n",
    "            \n",
    "            # 更新好叶子节点对的数量\n",
    "            for ld in left_depths:\n",
    "                for rd in right_depths:\n",
    "                    if ld + rd <= distance:\n",
    "                        self.count += 1\n",
    "                        \n",
    "            # 返回当前节点的所有叶子节点的深度，增加1，并且过滤掉超过 distance 的深度\n",
    "            return [d+1 for d in left_depths + right_depths if d+1 < distance]\n",
    "        \n",
    "        dfs(root)\n",
    "        return self.count\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def DFS(root):\n",
    "            if root == None:\n",
    "                return []\n",
    "            if root.left == None and root.right == None:\n",
    "                return [1]\n",
    "            left_depths = DFS(root.left)\n",
    "            right_depths = DFS(root.right)\n",
    "            for ld in left_depths:\n",
    "                for rd in right_depths:\n",
    "                    if ld + rd <= distance:\n",
    "                        nonlocal ans\n",
    "                        ans += 1\n",
    "            return [ld + 1 for ld in left_depths] + [rd + 1 for rd in right_depths]\n",
    "\n",
    "        ans = 0\n",
    "        DFS(root)\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        ans = 0\n",
    "        def helper(root):\n",
    "            nonlocal ans\n",
    "\n",
    "            if root is None:\n",
    "                return []\n",
    "            if root.left is None and root.right is None:\n",
    "                return [1]\n",
    "            \n",
    "            left_part = helper(root.left)\n",
    "            right_part = helper(root.right)\n",
    "\n",
    "            for n1 in left_part:\n",
    "                for n2 in right_part:\n",
    "                    if n1 + n2 <= distance:\n",
    "                        ans += 1\n",
    "            \n",
    "            cur = []\n",
    "            cur.extend(left_part)\n",
    "            cur.extend(right_part)\n",
    "            for i in range(len(cur)):\n",
    "                cur[i] += 1\n",
    "            return cur\n",
    "            \n",
    "        helper(root)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.count = 0\n",
    "        self.dfs(root, distance)\n",
    "        return self.count\n",
    "\n",
    "    def dfs(self, node, distance):\n",
    "        if not node:\n",
    "            return []\n",
    "\n",
    "        if not node.left and not node.right:\n",
    "            return [1]\n",
    "\n",
    "        left = self.dfs(node.left, distance)\n",
    "        right = self.dfs(node.right, distance)\n",
    " \n",
    "        for l in left:\n",
    "            for r in right:\n",
    "                if l + r <= distance:\n",
    "                    self.count += 1\n",
    "        \n",
    "        return [l+1 for l in left] + [r+1 for r in right]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        result = 0\n",
    "        \n",
    "        def dfs(root):\n",
    "            nonlocal result\n",
    "            if root is None:\n",
    "                return dict()\n",
    "            if root.left is None and root.right is None:\n",
    "                return {root: 1}\n",
    "            left = dfs(root.left)\n",
    "            right = dfs(root.right)\n",
    "            for v1 in left.values():\n",
    "                for v2 in right.values():\n",
    "                    if v1 + v2 <= distance:\n",
    "                        result += 1\n",
    "            left.update(right)\n",
    "            for k in left:\n",
    "                left[k] += 1\n",
    "            return left\n",
    "        \n",
    "        dfs(root)\n",
    "        return result\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 1024*1024 10^6次方的复杂度，过了再说\n",
    "        res = 0\n",
    "        def dfs(node):\n",
    "            nonlocal res\n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            L = []\n",
    "            R = []   \n",
    "            if node.left:\n",
    "                L = dfs(node.left)\n",
    "            if node.right:\n",
    "                R = dfs(node.right)\n",
    "            \n",
    "            for i in L:\n",
    "                for j in R:\n",
    "                    if i+j <= distance:\n",
    "                        res += 1\n",
    "            ans = L+R \n",
    "            for i in range(len(L+R)):\n",
    "                ans[i] += 1\n",
    "            return ans\n",
    "        dfs(root)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if not root or not root.left and not root.right: \n",
    "            return 0\n",
    "        \n",
    "        def _dfs(root):\n",
    "            \"\"\"获取左右子树中叶子结点到当前节点的距离\"\"\"\n",
    "            if not root: \n",
    "                return {}\n",
    "            # 叶子节点\n",
    "            if not root.left and not root.right: \n",
    "                return {root: 0}\n",
    "            \n",
    "            left_leaf = _dfs(root.left)\n",
    "            right_leaf = _dfs(root.right)\n",
    "            # 距离加1\n",
    "            for k, v in left_leaf.items(): left_leaf[k] = v + 1\n",
    "            for k, v in right_leaf.items(): right_leaf[k] = v + 1\n",
    "            \n",
    "            for lk, lv in left_leaf.items():\n",
    "                for rk, rv in right_leaf.items():\n",
    "                    if lv + rv <= distance:\n",
    "                        self.ans += 1\n",
    "            \n",
    "            # 合并左右子树的叶子节点，向上返回\n",
    "            left_leaf.update(right_leaf)\n",
    "            return left_leaf\n",
    "\n",
    "        self.ans = 0\n",
    "        _dfs(root)\n",
    "        return self.ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        \n",
    "        res = 0\n",
    "        def dfs(root):\n",
    "            nonlocal res\n",
    "            if root == None:\n",
    "                return []\n",
    "            \n",
    "            l_branch = dfs(root.left)\n",
    "            r_branch = dfs(root.right)\n",
    "\n",
    "            counts = [0]*11\n",
    "            for leftH, cntL in enumerate(l_branch):\n",
    "                for rightH, cntR in enumerate(r_branch):\n",
    "                    if leftH + rightH + 2> distance: break\n",
    "                    # print(root.val, leftH, rightH, cntL, cntR, cntL*cntR)\n",
    "                    res += cntL * cntR\n",
    "\n",
    "            if root.left is None and root.right is None:\n",
    "                counts[0] = 1\n",
    "                return counts\n",
    "\n",
    "\n",
    "            for i in range(1, 11):\n",
    "                if l_branch:\n",
    "                    counts[i] += l_branch[i-1]\n",
    "                if r_branch:\n",
    "                    counts[i] += r_branch[i-1] \n",
    "            return counts\n",
    "        dfs(root)\n",
    "        return res\n",
    "\n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.ans = 0\n",
    "        def dfs(node):\n",
    "            if not node:\n",
    "                return []\n",
    "            \n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            \n",
    "            left = dfs(node.left)\n",
    "            right = dfs(node.right)\n",
    "\n",
    "            for ii in left:\n",
    "                for jj in right:\n",
    "                    if ii + jj <= distance:\n",
    "                        self.ans += 1\n",
    "            \n",
    "            ret = left + right\n",
    "            ret = [(x + 1) for x in ret]\n",
    "            return ret\n",
    "        \n",
    "        dfs(root)\n",
    "        return self.ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        \n",
    "        res = 0\n",
    "        def dfs(root):\n",
    "            nonlocal res\n",
    "            if root == None:\n",
    "                return []\n",
    "            \n",
    "            l_branch = dfs(root.left)\n",
    "            r_branch = dfs(root.right)\n",
    "\n",
    "            counts = [0]*11\n",
    "            for leftH, cntL in enumerate(l_branch):\n",
    "                for rightH, cntR in enumerate(r_branch):\n",
    "                    if leftH + rightH + 2> distance: break\n",
    "                    # print(root.val, leftH, rightH, cntL, cntR, cntL*cntR)\n",
    "                    res += cntL * cntR\n",
    "\n",
    "            if root.left is None and root.right is None:\n",
    "                counts[0] = 1\n",
    "                return counts\n",
    "\n",
    "\n",
    "            for i in range(1, 11):\n",
    "                if l_branch:\n",
    "                    counts[i] += l_branch[i-1]\n",
    "                if r_branch:\n",
    "                    counts[i] += r_branch[i-1] \n",
    "            return counts\n",
    "        dfs(root)\n",
    "        return res\n",
    "\n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        \n",
    "        res = 0\n",
    "        def dfs(root):\n",
    "            nonlocal res\n",
    "            if root == None:\n",
    "                return []\n",
    "            \n",
    "            l_branch = dfs(root.left)\n",
    "            r_branch = dfs(root.right)\n",
    "\n",
    "            counts = [0]*11\n",
    "            for leftH, cntL in enumerate(l_branch):\n",
    "                for rightH, cntR in enumerate(r_branch):\n",
    "                    # if leftH + rightH + 2> distance: break\n",
    "                    if leftH + rightH > distance: break\n",
    "                    res += cntL * cntR\n",
    "\n",
    "            if root.left is None and root.right is None:\n",
    "                # counts[0] = 1\n",
    "                counts[1] = 1\n",
    "                return counts\n",
    "\n",
    "\n",
    "            for i in range(1, 11):\n",
    "                if l_branch:\n",
    "                    counts[i] += l_branch[i-1]\n",
    "                if r_branch:\n",
    "                    counts[i] += r_branch[i-1] \n",
    "            return counts\n",
    "        dfs(root)\n",
    "        return res\n",
    "\n",
    "                        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if not root or not root.left and not root.right: \n",
    "            return 0\n",
    "        \n",
    "        def _dfs(root):\n",
    "            \"\"\"获取左右子树中叶子结点到当前节点的距离\"\"\"\n",
    "            if not root: \n",
    "                return {}\n",
    "            # 叶子节点\n",
    "            if not root.left and not root.right: \n",
    "                return {root: 0}\n",
    "            \n",
    "            left_leaf = _dfs(root.left)\n",
    "            right_leaf = _dfs(root.right)\n",
    "            # 距离加1\n",
    "            for k, v in left_leaf.items(): left_leaf[k] = v + 1\n",
    "            for k, v in right_leaf.items(): right_leaf[k] = v + 1\n",
    "            \n",
    "            for lk, lv in left_leaf.items():\n",
    "                for rk, rv in right_leaf.items():\n",
    "                    if lv + rv <= distance:\n",
    "                        self.ans += 1\n",
    "            \n",
    "            # 合并左右子树的叶子节点，向上返回\n",
    "            left_leaf.update(right_leaf)\n",
    "            return left_leaf\n",
    "\n",
    "        self.ans = 0\n",
    "        _dfs(root)\n",
    "        return self.ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if distance==1:\n",
    "            return 0\n",
    "        ans=0\n",
    "        def dfs(r):\n",
    "            nonlocal ans \n",
    "            res=[0]*distance\n",
    "            if  not r.left and not r.right:\n",
    "                return True,res\n",
    "            if r.left:\n",
    "                y=dfs(r.left)\n",
    "                if y[0]:\n",
    "                    res[1]+=1\n",
    "                else:\n",
    "                    x=y[1]\n",
    "                    for i in range(distance-1):\n",
    "                        res[i+1]+=x[i]\n",
    "            if r.right:\n",
    "                y=dfs(r.right)\n",
    "                if y[0]:\n",
    "                    ans+=sum(res)\n",
    "                    res[1]+=1                    \n",
    "                else:\n",
    "                    x=y[1]\n",
    "                    for i in range(distance):\n",
    "                        ans+=sum(res[:distance-i])*x[i]\n",
    "                    for i in range(distance-1):\n",
    "                        res[i+1]+=x[i]\n",
    "            return False,res \n",
    "        dfs(root)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root):\n",
    "            if not root: return {}\n",
    "            if not root.left and not root.right: return {root:0}\n",
    "            l = dfs(root.left)\n",
    "            r = dfs(root.right)\n",
    "            for k in l: l[k] += 1\n",
    "            for k in r: r[k] += 1\n",
    "            for lv in l.values():\n",
    "                for rv in r.values():\n",
    "                    if lv + rv <= distance:\n",
    "                        nonlocal ans\n",
    "                        ans += 1\n",
    "            l.update(r)\n",
    "            return l \n",
    "        \n",
    "        ans = 0\n",
    "        dfs(root)\n",
    "        return ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        ans = 0\n",
    "        def dfs(node):\n",
    "            nonlocal ans\n",
    "            if not node:\n",
    "                return []\n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            left = dfs(node.left)\n",
    "            right = dfs(node.right)\n",
    "        \n",
    "            for l in left:\n",
    "                for r in right:\n",
    "                    ans += (l + r) <= distance\n",
    "            retval = left + right\n",
    "            for i in range(len(retval)):\n",
    "                retval[i] += 1\n",
    "            return retval\n",
    "        dfs(root)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def __init__(self):\n",
    "        self.res = 0\n",
    "\n",
    "    def dfs(self, root, distance):\n",
    "        if not root:\n",
    "            return {}\n",
    "        if not root.left and not root.right:\n",
    "            return {root: 0}\n",
    "        left = self.dfs(root.left, distance)\n",
    "        right = self.dfs(root.right, distance)\n",
    "        for k, v in left.items():\n",
    "            left[k] = v + 1\n",
    "        for k, v in right.items():\n",
    "            right[k] = v + 1\n",
    "        for lk,lv in left.items():\n",
    "            for rk,rv in right.items():\n",
    "                if lv+rv <= distance:\n",
    "                    self.res += 1\n",
    "        left.update(right)\n",
    "        return left\n",
    "    \n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if not root or (not root.left and not root.right):\n",
    "            return 0\n",
    "        self.dfs(root, distance)\n",
    "        return self.res\n",
    "        \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # return: (at depth i, #leaves (array), #pairs)\n",
    "        def dfs(node, distance):\n",
    "            dp = [0]*(distance+1)  # at depth i, #leaves\n",
    "            if not node.left and not node.right:\n",
    "                dp[0] = 1\n",
    "                return dp, 0\n",
    "            ldp, rdp = [0]*(distance+1), [0]*(distance+1)\n",
    "            lcnt = rcnt = 0\n",
    "\n",
    "            if node.left:\n",
    "                ldp, lcnt = dfs(node.left, distance)\n",
    "            if node.right:\n",
    "                rdp, rcnt = dfs(node.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                dp[i+1] += ldp[i]+rdp[i]\n",
    "            \n",
    "            cnt = 0  # use current node as parent, count #pairs\n",
    "            for i in range(distance+1):\n",
    "                for j in range(distance-i-1):\n",
    "                    cnt += ldp[i]*rdp[j]\n",
    "            return dp, cnt+lcnt+rcnt\n",
    "        \n",
    "        return dfs(root, distance)[1]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def f(node):\n",
    "            if not node:\n",
    "                return [], 0\n",
    "            nodes1, c1 = f(node.left)\n",
    "            nodes2, c2 = f(node.right)\n",
    "            c3 = 0\n",
    "            nodes3 = []\n",
    "            if not node.left and not node.right:\n",
    "                nodes3.append(0)\n",
    "            nodes3.extend([x+1 for x in nodes1])\n",
    "            nodes3.extend([x+1 for x in nodes2])\n",
    "            for x in nodes1:\n",
    "                for y in nodes2:\n",
    "                    if x+y+2 <= distance:\n",
    "                        c3 += 1\n",
    "            return nodes3, c1+c2+c3\n",
    "        x, c = f(root)\n",
    "        #print(x)\n",
    "        return c\n",
    "            \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        count = 0\n",
    "\n",
    "        def dfs(node): # return distances to leaf nodes\n",
    "            nonlocal count\n",
    "            if not node:\n",
    "                return []\n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            left = dfs(node.left)\n",
    "            right = dfs(node.right)\n",
    "            count += sum(l + r <= distance for l in left for r in right) # 在后序位置统计\n",
    "            return [n + 1 for n in left + right if n + 1 < distance] # 剪枝\n",
    "\n",
    "        dfs(root)\n",
    "        return count\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if not root or not root.left and not root.right: \n",
    "            return 0\n",
    "        \n",
    "        def _dfs(root):\n",
    "            \"\"\"获取左右子树中叶子结点到当前节点的距离\"\"\"\n",
    "            if not root: \n",
    "                return {}\n",
    "            # 叶子节点\n",
    "            if not root.left and not root.right: \n",
    "                return {root: 0}\n",
    "            \n",
    "            left_leaf = _dfs(root.left)\n",
    "            right_leaf = _dfs(root.right)\n",
    "            # 距离加1\n",
    "            for k, v in left_leaf.items(): left_leaf[k] = v + 1\n",
    "            for k, v in right_leaf.items(): right_leaf[k] = v + 1\n",
    "            \n",
    "            for lk, lv in left_leaf.items():\n",
    "                for rk, rv in right_leaf.items():\n",
    "                    if lv + rv <= distance:\n",
    "                        self.ans += 1\n",
    "            \n",
    "            # 合并左右子树的叶子节点，向上返回\n",
    "            left_leaf.update(right_leaf)\n",
    "            return left_leaf\n",
    "\n",
    "        self.ans = 0\n",
    "        _dfs(root)\n",
    "        return self.ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def solve(self,root,distance):\n",
    "        if(root is None):\n",
    "            return []\n",
    "        if(root.left==root.right):\n",
    "            return [1]\n",
    "        l_list=self.solve(root.left,distance)\n",
    "        r_list=self.solve(root.right,distance)\n",
    "        for l in l_list:\n",
    "            for r in r_list:\n",
    "                if(l+r<=distance):\n",
    "                    self.res+=1\n",
    "        return [item+1 for item in l_list+r_list]\n",
    "        \n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.res=0\n",
    "        self.solve(root,distance)\n",
    "        return self.res\n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def tree(self, root):\n",
    "        if not root: return []\n",
    "        if root.left is None and root.right is None: return [1]\n",
    "        l, r = self.tree(root.left), self.tree(root.right)\n",
    "        for i in l:\n",
    "            for j in r: \n",
    "                if i+j<=self.distance: self.res += 1\n",
    "        return [i+1 for i in l+r]\n",
    "\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.res, self.distance = 0, distance\n",
    "        self.tree(root)\n",
    "        return self.res\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.ans = 0\n",
    "        def dfs(node):\n",
    "            if not node:\n",
    "                return []\n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            l = dfs(node.left)\n",
    "            r = dfs(node.right)\n",
    "            self.ans += sum(1 for x in l for y in r if x + y <= distance)\n",
    "            return [x + 1 for x in l + r if x < distance]\n",
    "        dfs(root)\n",
    "        return self.ans\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.ans = 0\n",
    "        self.dfs(root, distance)\n",
    "        return self.ans\n",
    "    \n",
    "    def dfs(self, root, distance):\n",
    "        if not root:\n",
    "            return []\n",
    "\n",
    "        if root and not root.left and not root.right:\n",
    "            return [0]\n",
    "\n",
    "        l_heights = self.dfs(root.left, distance)\n",
    "        r_heights = self.dfs(root.right, distance)\n",
    "\n",
    "        l_heights = [h + 1 for h in l_heights]\n",
    "        r_heights = [h + 1 for h in r_heights]\n",
    "        \n",
    "        for i in range(len(l_heights)):\n",
    "            for j in range(len(r_heights)):\n",
    "                if l_heights[i] + r_heights[j] <= distance:\n",
    "                    self.ans += 1\n",
    "        \n",
    "        return l_heights + r_heights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, d: int) -> int:\n",
    "        if not root:\n",
    "            return 0\n",
    "        \n",
    "        def dfs(root):\n",
    "            nonlocal ans\n",
    "            if not (root.left or root.right):\n",
    "                return [1]\n",
    "            \n",
    "            l = dfs(root.left) if root.left else []\n",
    "            r = dfs(root.right) if root.right else []\n",
    "\n",
    "            for i, c1 in enumerate(l, start=1):\n",
    "                for c2 in r[:min(d-i, len(r))]:\n",
    "                    ans += c1 * c2\n",
    "            \n",
    "            cur = [0] * (max(len(l), len(r))+1)\n",
    "            for i in range(len(l)):\n",
    "                cur[i+1] += l[i]\n",
    "            for i in range(len(r)):\n",
    "                cur[i+1] += r[i]\n",
    "            \n",
    "            # 只返回距离小于distance的结点个数\n",
    "            return cur[:min(d-1, len(cur))]\n",
    "        \n",
    "        ans = 0\n",
    "        dfs(root)\n",
    "        return ans\n",
    "\n",
    "                    \n",
    "\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root, distance):\n",
    "            # index: distance from root\n",
    "            # value: number of leaves\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCnt = rightCnt = 0\n",
    "            if root.left:\n",
    "                leftDepths, leftCnt = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCnt = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                cnt += sum([leftDepths[i] * rightDepths[j] for j in range(distance - i - 1)])\n",
    "            \n",
    "            return (depths, cnt + leftCnt + rightCnt)\n",
    "        \n",
    "        _, ans = dfs(root, distance)\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        if not root or not root.left and not root.right: \n",
    "            return 0\n",
    "        \n",
    "        def _dfs(root):\n",
    "            \"\"\"获取左右子树中叶子结点到当前节点的距离\"\"\"\n",
    "            if not root: \n",
    "                return {}\n",
    "            # 叶子节点\n",
    "            if not root.left and not root.right: \n",
    "                return {root: 0}\n",
    "            \n",
    "            left_leaf = _dfs(root.left)\n",
    "            right_leaf = _dfs(root.right)\n",
    "            # 距离加1\n",
    "            for k, v in left_leaf.items(): left_leaf[k] = v + 1\n",
    "            for k, v in right_leaf.items(): right_leaf[k] = v + 1\n",
    "            \n",
    "            for lk, lv in left_leaf.items():\n",
    "                for rk, rv in right_leaf.items():\n",
    "                    if lv + rv <= distance:\n",
    "                        self.ans += 1\n",
    "            \n",
    "            # 合并左右子树的叶子节点，向上返回\n",
    "            left_leaf.update(right_leaf)\n",
    "            return left_leaf\n",
    "\n",
    "        self.ans = 0\n",
    "        _dfs(root)\n",
    "        return self.ans\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root, distance):\n",
    "            depths = [0]*(distance + 1)\n",
    "\n",
    "            left_depth = [0]*(distance + 1)\n",
    "            right_depth = [0]*(distance + 1)\n",
    "\n",
    "            left_count = 0\n",
    "            right_count = 0\n",
    "\n",
    "            if not root.left and not root.right:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            if root.left:\n",
    "                left_depth, left_count = dfs(root.left, distance)\n",
    "            \n",
    "            if root.right:\n",
    "                right_depth, right_count = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i+1] += left_depth[i]\n",
    "                depths[i+1] += right_depth[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += left_depth[i]*right_depth[j]\n",
    "            \n",
    "            return (depths, cnt + left_count + right_count)\n",
    "        \n",
    "        _, ret = dfs(root, distance)\n",
    "\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root):\n",
    "            if not root:\n",
    "                return {}\n",
    "            if not root.left and not root.right:\n",
    "                return {root: 0}\n",
    "            \n",
    "            lefts = dfs(root.left)\n",
    "            rights = dfs(root.right)\n",
    "            for leaf, dis in lefts.items():\n",
    "                lefts[leaf] = dis + 1\n",
    "            for leaf, dis in rights.items():\n",
    "                rights[leaf] = dis + 1\n",
    "            \n",
    "            nonlocal ans\n",
    "            for left, disLeft in lefts.items():\n",
    "                for right, disRight in rights.items():\n",
    "                    if disLeft + disRight <= distance:\n",
    "                        ans += 1\n",
    "\n",
    "            lefts.update(rights)\n",
    "            return lefts\n",
    "        \n",
    "        ans = 0\n",
    "        dfs(root)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        \n",
    "        self.res = 0\n",
    "\n",
    "        def dfs(node):\n",
    "            if not node:\n",
    "                return []\n",
    "            \n",
    "            if not node.left and not node.right:\n",
    "                return [1]\n",
    "            \n",
    "            l = dfs(node.left)\n",
    "            r = dfs(node.right)\n",
    "\n",
    "            for i in l:\n",
    "                for j in r:\n",
    "                    if i+j <= distance:\n",
    "                        self.res += 1\n",
    "            \n",
    "            return [x+1 for x in l+r]\n",
    "\n",
    "        dfs(root)\n",
    "        return self.res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(node):\n",
    "            if node is None:\n",
    "                return []\n",
    "            if node.left is None and node.right is None:\n",
    "                return [0]\n",
    "            left_res = dfs(node.left)\n",
    "            right_res = dfs(node.right)\n",
    "            filter_node = [node + 1 for node in left_res + right_res if node + 1 < distance]\n",
    "            for l_d in left_res:\n",
    "                for r_d in right_res:\n",
    "                    if l_d + r_d + 2 <= distance:\n",
    "                        res_list[0] += 1\n",
    "            return filter_node\n",
    "        res_list = [0]\n",
    "        dfs(root)\n",
    "        return res_list[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        self.ans = 0\n",
    "        def dfs(root) -> List[int]:\n",
    "            dist = [0] * 11\n",
    "            if root is None:\n",
    "                return dist\n",
    "            if root.left == root.right:\n",
    "                dist[0] = 1\n",
    "                return dist\n",
    "            l_dist, r_dist = dfs(root.left), dfs(root.right)\n",
    "            for i in range(11):\n",
    "                for j in range(11):\n",
    "                    if i + j + 2 <= distance:\n",
    "                        self.ans += l_dist[i] * r_dist[j]\n",
    "            for i in range(1, 11):\n",
    "                dist[i] = l_dist[i - 1] + r_dist[i - 1]\n",
    "            return dist\n",
    "        dfs(root)\n",
    "        return self.ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# https://leetcode.cn/problems/number-of-good-leaf-nodes-pairs/solutions/357905/hao-xie-zi-jie-dian-dui-de-shu-liang-by-leetcode-s/\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root, distance):\n",
    "            depths = [0] * (distance + 1)\n",
    "            is_leaf = not root.left and not root.right\n",
    "            if is_leaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount, rightCount = 0, 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return  (depths, cnt + leftCount + rightCount)\n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 对于 dfs(root,distance)，同时返回：\n",
    "        # 每个叶子节点与 root 之间的距离\n",
    "        # 以 root 为根节点的子树中好叶子节点对的数量\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # return: (each leaf distance,  #pairs)\n",
    "        def dfs(node, distance):\n",
    "            dp = [0]*(distance+1)\n",
    "            leaf = not node.left and not node.right\n",
    "            if leaf:\n",
    "                dp[0] = 1\n",
    "                return dp, 0\n",
    "            ldp, rdp = [0]*(distance+1), [0]*(distance+1)\n",
    "            lcnt = rcnt = 0\n",
    "\n",
    "            if node.left:\n",
    "                ldp, lcnt = dfs(node.left, distance)\n",
    "            if node.right:\n",
    "                rdp, rcnt = dfs(node.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                dp[i+1] += ldp[i]\n",
    "                dp[i+1] += rdp[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance+1):\n",
    "                for j in range(distance-i-1):\n",
    "                    cnt += ldp[i]*rdp[j]\n",
    "            return dp, cnt+lcnt+rcnt\n",
    "        \n",
    "        return dfs(root, distance)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root, distance):\n",
    "            # index: distance from root\n",
    "            # value: number of leaves\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCnt = rightCnt = 0\n",
    "            if root.left:\n",
    "                leftDepths, leftCnt = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCnt = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                cnt += sum([leftDepths[i] * rightDepths[j] for j in range(distance - i - 1)])\n",
    "            \n",
    "            return (depths, cnt + leftCnt + rightCnt)\n",
    "        \n",
    "        _, ans = dfs(root, distance)\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 对于 dfs(root,distance)，同时返回：\n",
    "        # 每个叶子节点与 root 之间的距离\n",
    "        # 以 root 为根节点的子树中好叶子节点对的数量\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 对于 dfs(root,distance)，同时返回：\n",
    "        # 每个叶子节点与 root 之间的距离\n",
    "        # 以 root 为根节点的子树中好叶子节点对的数量\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        ans = 0\n",
    "        def f(root):\n",
    "            nonlocal ans\n",
    "            a = defaultdict(int)\n",
    "            if not root.left and not root.right:\n",
    "                a[0] = 1\n",
    "                return a\n",
    "            l = defaultdict(int)\n",
    "            r = defaultdict(int)\n",
    "            if root.left:\n",
    "                l = f(root.left)\n",
    "                print(l)\n",
    "                for k in l:\n",
    "                    a[k + 1] = l[k]\n",
    "            if root.right:\n",
    "                r = f(root.right)\n",
    "                print(r)\n",
    "                for k in r:\n",
    "                    a[k + 1] += r[k]\n",
    "                    if l:\n",
    "                        for j in range(0, distance - k - 1):\n",
    "                            ans += r[k] * l[j]\n",
    "            return a\n",
    "        print(f(root))\n",
    "        return ans\n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        def dfs(root, distance) -> (List[int], int):\n",
    "            depth = [0] * (distance + 1)\n",
    "            count = 0\n",
    "            is_leaf = True if not root.left and not root.right else False\n",
    "            if is_leaf:\n",
    "                depth[0] = 1\n",
    "                return depth,0\n",
    "            left_depth, left_count = [0] * (distance + 1), 0\n",
    "            right_depth, right_count = [0] * (distance + 1), 0\n",
    "            if root.left:\n",
    "                left_depth, left_count = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                right_depth, right_count = dfs(root.right, distance)\n",
    "\n",
    "            for i in range(distance):\n",
    "                depth[i + 1] += left_depth[i]\n",
    "                depth[i + 1] += right_depth[i]\n",
    "\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    count += left_depth[i] * right_depth[j]\n",
    "            return depth, count + left_count + right_count\n",
    "\n",
    "        _, res = dfs(root, distance)\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 对于 dfs(root,distance)，同时返回：\n",
    "        # 每个叶子节点与 root 之间的距离\n",
    "        # 以 root 为根节点的子树中好叶子节点对的数量\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        # 对于 dfs(root,distance)，同时返回：\n",
    "        # 每个叶子节点与 root 之间的距离\n",
    "        # 以 root 为根节点的子树中好叶子节点对的数量\n",
    "        def dfs(root: TreeNode, distance: int) -> (List[int], int):\n",
    "            depths = [0] * (distance + 1)\n",
    "            isLeaf = not root.left and not root.right\n",
    "            if isLeaf:\n",
    "                depths[0] = 1\n",
    "                return (depths, 0)\n",
    "            \n",
    "            leftDepths, rightDepths = [0] * (distance + 1), [0] * (distance + 1)\n",
    "            leftCount = rightCount = 0\n",
    "\n",
    "            if root.left:\n",
    "                leftDepths, leftCount = dfs(root.left, distance)\n",
    "            if root.right:\n",
    "                rightDepths, rightCount = dfs(root.right, distance)\n",
    "            \n",
    "            for i in range(distance):\n",
    "                depths[i + 1] += leftDepths[i]\n",
    "                depths[i + 1] += rightDepths[i]\n",
    "            \n",
    "            cnt = 0\n",
    "            for i in range(distance + 1):\n",
    "                for j in range(distance - i - 1):\n",
    "                    cnt += leftDepths[i] * rightDepths[j]\n",
    "            \n",
    "            return (depths, cnt + leftCount + rightCount)\n",
    "        \n",
    "\n",
    "        _, ret = dfs(root, distance)\n",
    "        return ret\n",
    "\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "# Definition for a binary tree node.\n",
    "# class TreeNode:\n",
    "#     def __init__(self, val=0, left=None, right=None):\n",
    "#         self.val = val\n",
    "#         self.left = left\n",
    "#         self.right = right\n",
    "class Solution:\n",
    "    def countPairs(self, root: TreeNode, distance: int) -> int:\n",
    "        ans = 0\n",
    "        def f(root):\n",
    "            nonlocal ans\n",
    "            a = defaultdict(int)\n",
    "            if not root.left and not root.right:\n",
    "                a[0] = 1\n",
    "                return a\n",
    "            l = defaultdict(int)\n",
    "            r = defaultdict(int)\n",
    "            if root.left:\n",
    "                l = f(root.left)\n",
    "                for k in l:\n",
    "                    a[k + 1] = l[k]\n",
    "            if root.right:\n",
    "                r = f(root.right)\n",
    "                for k in r:\n",
    "                    a[k + 1] += r[k]\n",
    "                    if l:\n",
    "                        for j in range(0, distance - k - 1):\n",
    "                            ans += r[k] * l[j]\n",
    "            return a\n",
    "        f(root)\n",
    "        return ans\n",
    "        "
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
